import os
from PIL import Image
import cv2

def process_vision_info(messages):
    """
    处理聊天消息中的视觉信息，包括图像和视频。
    
    Args:
        messages: 包含用户和助手消息的列表，可能包含视觉内容
        
    Returns:
        image_inputs: 处理后的图像输入列表
        video_inputs: 处理后的视频输入列表
    """
    image_inputs = []
    video_inputs = []
    
    for message in messages:
        if message["role"] == "user":
            for content in message["content"]:
                if content["type"] == "image":
                    # 处理图像
                    image_path = content.get("image")
                    if image_path:
                        if os.path.exists(image_path):
                            image = Image.open(image_path)
                            image_inputs.append(image)
                
                elif content["type"] == "video":
                    # 处理视频列表
                    video_paths = content.get("video")
                    if isinstance(video_paths, list):
                        frames = []
                        for path in video_paths:
                            if os.path.exists(path):
                                image = Image.open(path)
                                frames.append(image)
                        if frames:
                            video_inputs.append(frames)
                    
                    # 处理单个视频文件
                    elif isinstance(video_paths, str) and os.path.exists(video_paths):
                        video_path = video_paths
                        cap = cv2.VideoCapture(video_path)
                        
                        # 获取视频总帧数
                        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                        
                        # 根据fps参数确定抽帧间隔
                        fps = content.get("fps", 1.0)
                        frame_interval = max(1, int(cap.get(cv2.CAP_PROP_FPS) / fps))
                        
                        # 抽取帧
                        frames = []
                        for i in range(0, total_frames, frame_interval):
                            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
                            ret, frame = cap.read()
                            if ret:
                                # 转换BGR到RGB并转为PIL Image
                                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                                pil_image = Image.fromarray(rgb_frame)
                                frames.append(pil_image)
                        
                        cap.release()
                        
                        if frames:
                            video_inputs.append(frames)
    
    return image_inputs, video_inputs 